#include "DesignExtractor.h"

#include "PKB.h"

#include <set>
#include <assert.h>
#include <algorithm>
#include <iostream>

DesignExtractor::DesignExtractor()
{
}

DesignExtractor::~DesignExtractor()
{
}

void DesignExtractor::ExtractDesign()
{
	// get list of procedures
	std::set<PROCINDEX> procedures = PKB::procTable.GetAllProcIndex();

	// for each procedure, traverse tree to extract relations
	for_each( procedures.begin(), procedures.end(), [ this ] ( PROCINDEX index )
	{
		TNODE* node = PKB::ast.GetRoot( index );
		assert( node && PKB::ast.GetNodeType( node ) == PROCEDURE );
		this->HandleProcedure( node );
	});
}

void DesignExtractor::HandleProcedure( TNODE* node )
{
	TNODE* stmtList = PKB::ast.GetChild( node, 0 );
	assert( stmtList );
	assert( PKB::ast.GetNodeType( stmtList ) == STMTLST );

	this->_procedure = node;

	this->HandleStatementList( stmtList );
}

void DesignExtractor::HandleStatementList( TNODE* node )
{
	TNODE* stmt = PKB::ast.GetChild( node, 0 );
	assert( stmt );
	assert( PKB::ast.GetNodeType( stmt ) == IF || PKB::ast.GetNodeType( stmt ) == WHILE || PKB::ast.GetNodeType( stmt ) == ASSIGN );

	this->HandleStatement( stmt );
}

void DesignExtractor::HandleStatement( TNODE* node )
{
	// Parent Relationship
	TNODE* parent = NULL;
	if ( this->_parent.size() > 0 )
		parent = this->_parent.at( 0 );

	if ( parent )
	{
		PKB::parent.InsertParent( this->STMTFromTNode( parent ), this->STMTFromTNode( node ) );
	}

	switch ( PKB::ast.GetNodeType( node ) )
	{
		case IF:
			this->HandleIf( node );
			break;

		case WHILE:
			this->HandleWhile( node );
			break;

		case CALL:
			this->HandleCall( node );
			break;

		case ASSIGN:
			this->HandleAssign( node );
			break;
	}

	TNODE* right_sibling = PKB::ast.GetRightSibling( node );
	if ( right_sibling )
	{
		if ( PKB::ast.GetNodeType( node ) != IF )
		{
			// Next Relationship
			PKB::next.InsertNext( PKB::ast.GetStatementNumber( node ), PKB::ast.GetStatementNumber( right_sibling ) );
		}

		// Follows Relationship
		PKB::follows.InsertFollows( this->STMTFromTNode( node ), this->STMTFromTNode( right_sibling ) );
		this->HandleStatement( right_sibling );
	}
	else
	{
		// Next Relationship
		// no right_sibling and has parent must link back for Next relationships
		// while: link back to parent node
		// if:    link back to parent's right_sibling. terminate when while is found or when if has right_sibling. ( must check all parents of parents )
		if ( parent )
		{
			for ( auto p = this->_parent.begin(); p != this->_parent.end(); p++ )
			{
				NODETYPE parent_type = PKB::ast.GetNodeType( *p );
				if ( parent_type == WHILE )
				{
					PKB::next.InsertNext( PKB::ast.GetStatementNumber( node ), PKB::ast.GetStatementNumber( *p ) );
					break;
				}
				else if ( parent_type == IF )
				{
					TNODE* parent_right_sibling = PKB::ast.GetRightSibling( *p );

					if ( !parent_right_sibling )
						continue;

					PKB::next.InsertNext( PKB::ast.GetStatementNumber( node ), PKB::ast.GetStatementNumber( parent_right_sibling ) );
					break;
				}
			}
		}
	}
}


void DesignExtractor::HandleIf( TNODE* node )
{
	TNODE* variable = PKB::ast.GetChild( node, 0 );
	TNODE* if_stmtList = PKB::ast.GetChild( node, 1 );
	TNODE* then_stmtList = PKB::ast.GetChild( node, 2 );
	assert( variable && if_stmtList && then_stmtList );
	
	// Next
	PKB::next.InsertNext( PKB::ast.GetStatementNumber( node ), PKB::ast.GetStatementNumber( PKB::ast.GetChild( if_stmtList, 0 ) ) );
	PKB::next.InsertNext( PKB::ast.GetStatementNumber( node ), PKB::ast.GetStatementNumber( PKB::ast.GetChild( then_stmtList, 0 ) ) );

	// ShortcutTable
	PKB::shortcutTable.InsertIf( PKB::ast.GetStatementNumber( node ) );

	STRING variable_name = PKB::ast.GetValue( variable );
	VARINDEX variable_index = PKB::varTable.GetVarIndex( variable_name );
	assert( variable_index >= 0 );

	// InsertUses for container and procedure
	STRING proc_name = PKB::ast.GetValue( this->_procedure );
	PROCINDEX proc_index = PKB::procTable.GetProcIndex( proc_name );

	STMT node_statement = this->STMTFromTNode( node );
	PKB::uses.InsertUses( node_statement, variable_index, proc_index );

	this->_parent.insert( this->_parent.begin(), node );

	this->HandleStatementList( if_stmtList );
	this->HandleStatementList( then_stmtList );

	this->_parent.erase( this->_parent.begin() );
}


void DesignExtractor::HandleWhile( TNODE* node )
{
	TNODE* variable = PKB::ast.GetChild( node, 0 );
	TNODE* stmtList = PKB::ast.GetChild( node, 1 );
	assert( variable && stmtList );

	// Next
	PKB::next.InsertNext( PKB::ast.GetStatementNumber( node ), PKB::ast.GetStatementNumber( PKB::ast.GetChild( stmtList, 0 ) ) );

	// ShortcutTable
	PKB::shortcutTable.InsertWhile( PKB::ast.GetStatementNumber( node ) );

	STRING variable_name = PKB::ast.GetValue( variable );
	VARINDEX variable_index = PKB::varTable.GetVarIndex( variable_name );
	assert( variable_index >= 0 );

	// Insert uses for container and procedure
	STRING proc_name = PKB::ast.GetValue( this->_procedure );
	PROCINDEX proc_index = PKB::procTable.GetProcIndex( proc_name );

	STMT node_statement = this->STMTFromTNode( node );
	PKB::uses.InsertUses( node_statement, variable_index, proc_index );

	if ( this->_parent.size() > 0 )
	{
		for ( auto it = this->_parent.begin(); it != this->_parent.end(); it++ )
		{
			TNODE* parent = *it;
			if ( parent )
			{
				PKB::uses.InsertUses( this->STMTFromTNode( parent ), variable_index, proc_index );
			}
		}
	}

	this->_parent.insert( this->_parent.begin(), node );

	this->HandleStatementList( stmtList );

	this->_parent.erase( this->_parent.begin() );
}

void DesignExtractor::HandleCall( TNODE* node )
{
	// Calls
	STRING proc_name = PKB::ast.GetValue( this->_procedure );
	PROCINDEX caller_index = PKB::procTable.GetProcIndex( proc_name );
	PROCINDEX callee_index = PKB::procTable.GetProcIndex( node->value );

	PKB::calls.InsertCalls( PKB::ast.GetStatementNumber( node ), caller_index, callee_index );

	// ShortcutTable
	PKB::shortcutTable.InsertCall( PKB::ast.GetStatementNumber( node ) );
}

void DesignExtractor::HandleAssign( TNODE* node )
{
	TNODE* variable = PKB::ast.GetChild( node, 0 );
	TNODE* expression = PKB::ast.GetChild( node, 1 );
	//assert( variable && expression );

	// ShortcutTable
	PKB::shortcutTable.InsertAssign( PKB::ast.GetStatementNumber( node ) );

	STRING variable_name = PKB::ast.GetValue( variable );
	VARINDEX variable_index = PKB::varTable.GetVarIndex( variable_name );
	assert( variable_index >= 0 );

	// Used constants
	INTEGER_ARRAY constants = this->FindConstantsInExpressionTree( expression );
	for_each( constants.begin(), constants.end(), [ this ] ( INTEGER value )
	{
		PKB::shortcutTable.InsertConstant( value );
	});

	// Used variables
	VARINDEX_ARRAY variables = this->FindVariablesInExpressionTree( expression );

	STRING proc_name = PKB::ast.GetValue( this->_procedure );
	PROCINDEX proc_index = PKB::procTable.GetProcIndex( proc_name );

	if ( this->_parent.size() > 0 )
	{
		for ( auto it = this->_parent.begin(); it != this->_parent.end(); it++ )
		{
			TNODE* parent = *it;
			if ( parent )
			{
				STMT parent_statement = this->STMTFromTNode( parent );
				for_each( variables.begin(), variables.end(), [ this, parent_statement, proc_index ] ( VARINDEX index )
				{
					PKB::uses.InsertUses( parent_statement, index, proc_index );
				});
				PKB::modifies.InsertModifies( parent_statement, variable_index, proc_index );
			}
		}
	}
	
	STMT node_statement = this->STMTFromTNode( node );
	for_each( variables.begin(), variables.end(), [ this, node_statement, proc_index ] ( VARINDEX index )
	{
		PKB::uses.InsertUses( node_statement, index, proc_index );
	});
	PKB::modifies.InsertModifies( node_statement, variable_index, proc_index );
}

VARINDEX_ARRAY DesignExtractor::FindVariablesInExpressionTree( TNODE* node )
{
	// recurse and collect variables
	return this->FindVariablesInExpressionTreeHelper( node );
}

VARINDEX_ARRAY DesignExtractor::FindVariablesInExpressionTreeHelper( TNODE* node )
{
	VARINDEX_ARRAY temp;
	VARINDEX_ARRAY variables;

	if ( PKB::ast.GetNodeType( node ) == VARIABLE )
	{
		STRING var_name = PKB::ast.GetValue( node );
		VARINDEX var_index = PKB::varTable.GetVarIndex( var_name );
		
		assert( var_index >= 0 );

		variables.insert( var_index );
	}

	// Children
	TNODE* left = PKB::ast.GetChild( node, 0 );
	TNODE* right = PKB::ast.GetChild( node, 1 );
	
	if ( left )
	{
		temp = this->FindVariablesInExpressionTreeHelper( left );
		variables.insert( temp.begin(), temp.end() );
	}

	if ( right )
	{
		temp = this->FindVariablesInExpressionTreeHelper( right );
		variables.insert( temp.begin(), temp.end() );
	}

	return variables;
}

INTEGER_ARRAY DesignExtractor::FindConstantsInExpressionTree( TNODE* node )
{
	// recurse and collect constants
	return this->FindConstantsInExpressionTreeHelper( node );
}

INTEGER_ARRAY DesignExtractor::FindConstantsInExpressionTreeHelper( TNODE* node )
{
	INTEGER_ARRAY temp;
	INTEGER_ARRAY constants;

	if ( PKB::ast.GetNodeType( node ) == CONSTANT )
	{
		STRING value = PKB::ast.GetValue( node );
		constants.insert( atoi( value.c_str() ) );
	}

	// Children
	TNODE* left = PKB::ast.GetChild( node, 0 );
	TNODE* right = PKB::ast.GetChild( node, 1 );
	
	if ( left )
	{
		temp = this->FindConstantsInExpressionTreeHelper( left );
		constants.insert( temp.begin(), temp.end() );
	}

	if ( right )
	{
		temp = this->FindConstantsInExpressionTreeHelper( right );
		constants.insert( temp.begin(), temp.end() );
	}

	return constants;
}

STMT DesignExtractor::STMTFromTNode( TNODE* node )
{
	STMT statement;
	statement.node_type = PKB::ast.GetNodeType( node );
	statement.statement_number = PKB::ast.GetStatementNumber( node );
	return statement;
}